import re
from typing import Union, List

import dspy

from .savetemp import savetemp


class GenPersona(dspy.Signature):
    """
    You need to select a group of academic experts who specialize in systematic research about this topic
    Each expert must represent a distinct perspective or role relevant to the topic, ensuring a comprehensive and balanced view. 
    The experts should have diverse academic or research-focused backgrounds directly or indirectly related to the topic, avoiding overly general perspectives.

    Provide the answer in the format: 
    1. short summary of expert 1: description\n2. short summary of expert 2: description\n...

    """

    topic = dspy.InputField(prefix="Topic of interest:", format=str)
    max_num_persona = dspy.InputField(prefix="Maximum number of personas to generate:", format=str)
    personas = dspy.OutputField(format=str)



class CreateWriterWithPersona(dspy.Module):
    """Discover different perspectives of researching the topic by reading Wikipedia pages of related topics."""

    def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]):
        super().__init__()
        
        self.gen_persona = dspy.ChainOfThought(GenPersona)
        self.engine = engine

    def forward(self, topic: str,max_num_persona: int, draft=None):
        with dspy.settings.context(lm=self.engine):

              gen_persona_output = self.gen_persona(
                topic=topic,max_num_persona=max_num_persona
              ).personas

        # savetemp
        savetemp("persona_generator/gen_persona_output.txt", gen_persona_output)

        # Extract personas from gen_persona_output
        personas = []
        for s in gen_persona_output.split("\n"):
            match = re.search(r"\d+\.\s*(.*)", s)
            if match:
                personas.append(match.group(1))


        return dspy.Prediction(
            personas=personas,
            raw_personas_output=personas,
        )



class StormPersonaGenerator:
    """
    A generator class for creating personas based on a given topic.

    This class uses an underlying engine to generate personas tailored to the specified topic.
    The generator integrates with a `CreateWriterWithPersona` instance to create diverse personas,
    including a default 'Basic fact writer' persona.

    Attributes:
        create_writer_with_persona (CreateWriterWithPersona): An instance responsible for
            generating personas based on the provided engine and topic.

    Args:
        engine (Union[dspy.dsp.LM, dspy.dsp.HFModel]): The underlying engine used for generating
            personas. It must be an instance of either `dspy.dsp.LM` or `dspy.dsp.HFModel`.
    """

    def __init__(self, engine: Union[dspy.dsp.LM, dspy.dsp.HFModel]):
        self.create_writer_with_persona = CreateWriterWithPersona(engine=engine)

    def generate_persona(self, topic: str, max_num_persona: int = 3) -> List[str]:
        """
        Generates a list of personas based on the provided topic, up to a maximum number specified.

        This method first creates personas using the underlying `create_writer_with_persona` instance
        and then prepends a default 'Basic fact writer' persona to the list before returning it.
        The number of personas returned is limited to `max_num_persona`, excluding the default persona.

        Args:
            topic (str): The topic for which personas are to be generated.
            max_num_persona (int): The maximum number of personas to generate, excluding the
                default 'Basic fact writer' persona.

        Returns:
            List[str]: A list of persona descriptions, including the default 'Basic fact writer' persona
                and up to `max_num_persona` additional personas generated based on the topic.
        """
        #向CreateWriterWithPersona中传递进入max_num_persona参数
        personas = self.create_writer_with_persona(topic=topic, max_num_persona=max_num_persona)
        default_persona = "Comprehensive Review Writer: Expert in integrating diverse methodologies to thoroughly analyze technologies or methods, avoiding superficial summaries."
        considered_personas = [default_persona] + personas.personas[:max_num_persona]
        return considered_personas
